// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief	Implementation of Logistic Model (Design Matrix)
*/
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "bitstring.h"
#include "bmapfile.h"
#include "chromosomes.h"
#include "logistic.h"
#include "pmaths.h"
#include "sample.h"
#include "utils.h"

void logistic_model_init(struct logistic_model *p) {
	p->miss = NULL;
	p->nMiss = 0;
	p->np = 0;
	p->all_valid = FALSE;

	// Flag variables associated with the design matrix of the logistic model.
	p->mAA = 0;
	p->mAB = 1;
	p->mBB = 2;

	// The n-D data arrays
	p->dimension = 0;
	p->additive = NULL;
	p->haploid = NULL;
	p->xchr = NULL;
	p->nInteraction = 1;

	// 2-D array X (indpendent array)
	p->X = NULL;
	p->nX = 0;
	p->X_max_used = 0;

	// 1-D dependent array.
	p->nind = 0;
	p->p = p->V = NULL;
	p->Y = NULL;

	// FIT LM variables
	p->coef = p->T3 = p->ncoef = NULL;
	p->S = p->T2 = p->Xt = NULL;
	p->S_inverse = NULL;
	p->T = p->T_inverse = NULL;

	p->valid = NULL;
}

int logistic_model_free_attributes(struct logistic_model *lm) {
	if (lm == NULL) return 0;
	else {
		if (lm->miss != NULL) {
			free(lm->miss);
			lm->miss = NULL;
			lm->nMiss = 0;
		}

		if (lm->valid != NULL) {
			free(lm->valid);
			lm->valid = NULL;
		}

		// Free the n-D results array.
		if (lm->dimension != 0) {
			if (lm->additive != NULL) free(lm->additive);
			if (lm->haploid != NULL) free(lm->haploid);
			if (lm->xchr != NULL) free(lm->xchr);
			lm->additive = NULL;
			lm->xchr = lm->haploid = NULL;
		}

		// Clean-up the 'X' design matrix component.
		if (lm->X != NULL) {
			free_matrix(&lm->X);
			lm->X = NULL;
			lm->nX = 0;
			lm->X_max_used = 0;
		}
		
		// Clean-up 'Y' bit of the design matrix
		if (lm->p != NULL) free(lm->p);
		if (lm->V != NULL) free(lm->V);
		if (lm->Y != NULL) free(lm->Y);
		lm->p = lm->V = NULL;
		lm->Y = NULL;

		// Clean up the FIT LM variables.
		if (lm->coef != NULL) free(lm->coef);
		if (lm->S != NULL) free_matrix(&lm->S);
		if (lm->S_inverse != NULL) free_matrix(&lm->S_inverse);
		if (lm->T2 != NULL) free_matrix(&lm->T2);
		if (lm->T3 != NULL) free(lm->T3);
		if (lm->ncoef != NULL) free(lm->ncoef);
		if (lm->Xt != NULL) free_matrix(&lm->Xt);
		if (lm->T != NULL) free_matrix(&lm->T);
		if (lm->T_inverse != NULL) free_matrix(&lm->T_inverse);
		lm->coef = lm->ncoef = NULL;
		lm->S = lm->T2 = lm->T = lm->T_inverse = NULL;

		if (lm->valid != NULL) {
			free(lm->valid);
			lm->valid = NULL;
		}

		// Clear the remaining index variables.
		lm->np = 0;
		lm->nind = 0;
	}

	return 1;
}

void logistic_model_set_np(struct logistic_model *lm, BOOL sex_effect) {
	// Number of parameters
	// Epistasis 1 for intercept column, 2 for additive effect considered and 1 for the interaction.
	lm->np = 1 + 2 + 1;

	// Sex effect?
	if (sex_effect) lm->np++;
}

int logistic_model_alloc(struct logistic_model *lm, int nSamples) {
	int dimension = 2;

	if (lm == NULL) return 0;
	else if (lm->np == 0) return 0;

	lm->coef = create_vector(lm->np, 0);
	lm->S = create_matrix(lm->np, lm->np, 0);
	lm->S_inverse = create_matrix(lm->np, lm->np, 0);
	lm->T2 = create_matrix(lm->np, nSamples, 0);
	lm->T3 = create_vector(nSamples, 0);
	lm->ncoef = create_vector(lm->np, 0);
	lm->Xt = create_matrix(lm->np, nSamples, 0);
	lm->T = create_matrix(lm->np, lm->np, 0);
	lm->T_inverse = create_matrix(lm->np, lm->np, 0);
	lm->additive = (int*) malloc(sizeof(int) * dimension);
	lm->haploid = (BOOL*) malloc(sizeof(BOOL) * dimension);
	lm->xchr = (BOOL*) malloc(sizeof(BOOL) * dimension);
	lm->dimension = dimension; // The n-wise character of the algorithm, at the moment N=2.
	lm->miss = (BOOL*) malloc(sizeof(BOOL) * nSamples);
	lm->nMiss = nSamples;
	lm->X = create_matrix(nSamples, lm->np, 0);
	lm->Y = (int*) malloc(sizeof(int) * nSamples);
	lm->p = create_vector(nSamples, 0);
	lm->V = create_vector(nSamples, 0);
	lm->valid = (BOOL*) malloc(sizeof(BOOL) * lm->np);
	
	if (lm->coef == NULL || lm->S == NULL || lm->S_inverse == NULL || lm->T2 == NULL ||
		lm->T3 == NULL || lm->ncoef == NULL || lm->Xt == NULL || 
		lm->T == NULL || lm->T_inverse == NULL || lm->additive == NULL || 
		lm->haploid == NULL ||  lm->xchr == NULL || lm->miss == NULL || lm->X == NULL ||
		lm->Y == NULL || lm->p == NULL || lm->V == NULL || lm->valid == NULL) 
		return 0; 
	
	return 1;
}

int logistic_model_add_additive_snp(struct logistic_model *lm, struct chromosomes *chromos, struct small_locus *m, int snp_index, int dimension_index) {
	if (lm == NULL || chromos == NULL || m == NULL || snp_index < 0 || dimension_index < 0)
		return 0;

	lm->additive[dimension_index] = snp_index;
	
	if (chromos->chr_sex[m[snp_index].chr]) 
		lm->xchr[dimension_index] = TRUE;
	else
		lm->xchr[dimension_index] = FALSE;

	if (chromos->chr_haploid[m[snp_index].chr]) 
		lm->haploid[dimension_index] = TRUE;
	else
		lm->haploid[dimension_index] = FALSE;

	return 1;
}

/*
Build X matrix (including intercept)
Iterate a person at a time, entering only 
valid rows into X (non-missing); also build Y
at the same time.
*/
enum MATHS_RETURN logistic_model_build_design_matrix(struct logistic_model *lm, char *ipedData, struct sample *samples, int nSamples, int nIpedRecordSizeInBytes, struct selected_options *ops) {
	int i, j;
	MATRIX_T X = NULL;  // Pointer linear model X co-ordinate.
	char *p = NULL;
	int ped_offset = nIpedRecordSizeInBytes * sizeof(char);
	
	if (lm == NULL || ipedData == NULL) return MEMORY_FAILURE;
	
	// Count the number of missing samples in the sample array (if any).
	// to get dimension for the 2-D vector array.
	for (i = 0, X = lm->X, lm->X_max_used = 0, p = ipedData; i < nSamples; i++, p += ped_offset) {
		struct sample *person = &samples[i];
		VECTOR_T trow = NULL; // Current row in the design matrix.
		int c = 1; 
		BOOL skip = FALSE;
		int *s = NULL; // 2-D list of SNP indexes to be considered for pairwise association things.
		int snp;
		
		// Ignore if missing phenotype, or the user set this to missing
		if (lm->miss[i]) continue;
		
		// List of indexes in the design matrix.
		// 0) Intercept      
		// 1) Main effects of SNPs 
		// 2) Interactions of the above

		trow = X[lm->X_max_used];
		trow[0] = 1; // Intercept

		// Additive effects (assuming individual-major mode)
		for (j = 0, s = lm->additive, snp = 0; j < lm->dimension; j++, s++, snp++) {
			long bitpos = (*s) * 2;
			unsigned int mask = 0x1; 
			BitString bs = (BitString) p;
			BOOL i1, i2;

			// Get first bit.
			mask = mask << (bitpos % BITBULKSIZE);
			mask = bs[bitpos/BITBULKSIZE] & mask;
			i1 = (mask !=0) ? TRUE : FALSE;
			
			// Get second bit
			bitpos++;
			mask = 0x1; 
			mask = mask << (bitpos % BITBULKSIZE);
			mask = bs[bitpos/BITBULKSIZE] & mask;
			i2 = (mask !=0) ? TRUE : FALSE;

			//if (snp == 0) printf("%i %i %i %i\n", *s, snp, i1, i2);

			if (lm->xchr[snp]) {
				// X chromosome coding
				double mA = 0, mB = 1;
				if (ops->xchr_model == 1) {
					mA = 0.5; 
					mB = 1.5;
				} else if ( ops->xchr_model == 2 ) {
					mA = 0;
					mB = 2;
				} else if ( ops->xchr_model > 2 ) {
					mA = 0;
					mB = 1;
				}

				// male
				if ( person->sex ) {
					if (i1) {
						if (!i2) {
							lm->miss[i] = TRUE;
							skip = TRUE;
						}
						else
							trow[c++] = mA; 	      
					} else {
						if (i2) {
							// This should not happen...
							lm->miss[i] = TRUE;
							skip = TRUE;
						}		  
						else
							trow[c++] = mB; 
					}
				} 
				// female x-chromosome
				else {
					if (i1) {
						if (!i2) {
							lm->miss[i] = TRUE;
							skip = TRUE;
						}
						else
							trow[c++] = 0; 	      
					} else {
						if (i2)
							trow[c++] = 1; // het 
						else
							trow[c++] = 2; // hom
					}		  
				}

			}
			else if (lm->haploid[snp]) {
				// Haploid coding
				if (i1) {
					if (!i2) {
						lm->miss[i] = TRUE;
						skip = TRUE;
					} else
						trow[c++] = 0; 	      
				} else {
					if (i2) {
						// This should not happen...
						lm->miss[i] = TRUE;
						skip = TRUE;
					}		  
					else
						trow[c++] = 1; 
				}
			} else {
				// Autosomal coding
				if (i1) {
					if (!i2) {
						lm->miss[i] = TRUE;
						skip = TRUE;
					}
					else
						trow[c++] = lm->mAA;
				}
				else {
					if (i2)
						trow[c++] = lm->mAB; // het 
					else
						trow[c++] = lm->mBB; // hom
				}
			} 
		}
		
		// Flag if need to ignore a row in the 2-D array.
		if (skip) continue; 
		
		// Sex effect (automatically set for X chromosome models)
		// Seems to be associated with GLM analysis only.
		if ( ops->sex_effect ) {
			if (samples[i].sex) trow[c++] = 1;
			else trow[c++] = 0;
		}

		// Interactions 
		for (j = 0; j < lm->nInteraction; j++)
			trow[c++] = trow[lm->interaction[j].p1] * trow[lm->interaction[j].p2];
		
		lm->X_max_used++; // If reach here, all is well so increment flag to use another row in the 'X' 2-D array.
	}
	//echo_matrix(lm->X, lm->X_max_used, lm->np);
	// VIF-based check for multicollinearity
	lm->all_valid = check_vif(lm->X, lm->X_max_used, lm->np);

	// Add Y variable also
	if(!logistic_model_set_dependent(lm, samples, nSamples)) return MEMORY_FAILURE; 
	
	return MATHS_SUCCESS;
}

// Set phenotype to 'aff' variable
int logistic_model_set_dependent(struct logistic_model *lm, struct sample *samples, int nSamples) {
	int i;

	if (lm == NULL || samples == NULL || lm->Y == NULL) return 0;
	else if (nSamples == 0) return 0;

	lm->nind = 0; // No. affected individuals.
	for (i = 0 ; i < nSamples; i++) {
		if (!lm->miss[i]) {
			if (samples[i].pperson->aff) lm->Y[lm->nind++] = 1;
			else lm->Y[lm->nind++] = 0;
		}
	}

	return 1;
}

enum MATHS_RETURN logistic_model_fitLM(struct logistic_model *lm) {
	BOOL converge = FALSE;
    int it = 0;
	BOOL callFailed = FALSE;
	int i, j, k;

	if (lm == NULL) return MATHS_FAILURE;
	else if (lm->np == 0 || lm->nind == 0 || !lm->all_valid) return MATHS_FAILURE;
   
	zero_vector(lm->coef, lm->np);
	
	// Newton-Raphson to fit logistic model
	zero_vector(lm->coef, lm->np);
	zero_matrix(lm->S, lm->np, lm->np);
	while (!converge && it < 20) {
		double delta = 0;
		
		// Determine p and V
		for (i = 0;  i< lm->nind; i++) {
			double t = 0;
			for (j = 0; j < lm->np; j++) t += lm->coef[j] * lm->X[i][j];	    
			lm->p[i] = 1/(1+exp(-t));
			lm->V[i] = lm->p[i] * (1-lm->p[i]);
		}

		for (j = 0; j < lm->np; j++) {
			for (k = j; k < lm->np; k++) {
				double sum = 0;
				for (i = 0; i< lm->nind; i++)
					sum += lm->X[i][j] * lm->V[i] * lm->X[i][k] ;
				lm->T[j][k] = lm->T[k][j] = sum;	      
			}
		}

		if (!svd_inverse(lm->T, lm->T_inverse, lm->np, lm->np)) {
			callFailed = TRUE;
			break;
		}
		
		// note implicit transpose of X
		zero_matrix(lm->T2, lm->np, lm->nind);
		for (i = 0; i < lm->np; i++)
			for (j = 0; j < lm->nind; j++)
				for (k = 0; k < lm->np; k++) lm->T2[i][j] += lm->T_inverse[i][k] * lm->X[j][k];

		for (i = 0; i < lm->nind; i++) 
			lm->T3[i] = lm->Y[i] - lm->p[i];

		zero_vector(lm->ncoef, lm->np);
		for (j = 0; j < lm->np; j++) 
			for (i = 0; i < lm->nind; i++) 
				lm->ncoef[j] += lm->T2[j][i] * lm->T3[i];

		// Update coefficients, and check for 
		// convergence
		for (j = 0; j < lm->np; j++) {
			delta += fabs(lm->ncoef[j]);
			lm->coef[j] += lm->ncoef[j];
		}
		if (delta < 1e-6) converge = TRUE;

		it++; // Next iteration
	}

	// Obtain covariance matrix of estimates
	// S <- solve( t(X) %*% V %*% X ) 
	// Transpose X and multiple by diagonal V
	for (i = 0; i < lm->nind; i++)
		for (j = 0; j < lm->np; j++) lm->Xt[j][i] = lm->X[i][j] * lm->V[i];

	mult_matrix(lm->Xt,lm->np,lm->nind,lm->X,lm->X_max_used,lm->np,lm->S);
	if (!svd_inverse(lm->S, lm->S_inverse, lm->np, lm->np)) {
		callFailed = TRUE; 
	} else {
		// Swap the S matrix for its inverse.
		MATRIX_T tmp = NULL;
		
		tmp = lm->S;
		lm->S = lm->S_inverse;
		lm->S_inverse =  tmp;	
	}
	
	if (callFailed) return MATHS_FAILURE;
	else return MATHS_SUCCESS; 
}

enum MATHS_RETURN logistic_model_validParameters(struct logistic_model *lm) {
	int i, j;

	if (lm == NULL) return MATHS_FAILURE;
	// Empty model?
	else if (lm->np == 0 || lm->nind == 0) {
		lm->all_valid = FALSE;
		return MATHS_FAILURE;
	}

	// Check for multicollinearity
	// For each term, see that estimate is not too strongly (r>0.99)
	// correlated with another, starting at last
	for (i = 1; i < lm->np; i++) {
		lm->valid[i] = TRUE;
		if ( lm->S[i][i] < 1e-20 ) 
			lm->valid[i] = lm->all_valid = FALSE; 
		else if ( ! realnum(lm->S[i][i]) ) 
			lm->valid[i] = lm->all_valid = FALSE; 
	}

	if (lm->all_valid) {
		for (i = lm->np - 1; i > 0; i--) {
			for (j = i-1; j >= 0; j--) {	  
				if (lm->S[i][j] / sqrt( lm->S[i][i] * lm->S[j][j] ) > 0.99999) {
					lm->valid[i] = FALSE;
					lm->all_valid = FALSE;
					break;
				}
			}
		}
	} 

	return MATHS_SUCCESS;
}

double logistic_model_get_statistic(struct logistic_model *lm, int testParameter) {
	if (lm->all_valid) {
		return ( lm->coef[testParameter] * lm->coef[testParameter] ) 
			/ lm->S[testParameter][testParameter];
	}
	else 
		return 0;
}